# -*- coding: utf-8 -*-
"""
Created on Fri May 12 10:16:24 2023

@author: lenovo
"""

import os
import numpy as np
from osgeo import gdal
import glob
# import shutil
# from tqdm import tqdm # 运行计时
import matplotlib.pyplot as plt
################################################
def read_img(filename):
    dataset = gdal.Open(filename)  # 打开文件
    if dataset == None:
        print(filename+"文件无法打开")
        return
    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    # im_bands = dataset.RasterCount #波段数
    # im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
    # im_proj = dataset.GetProjection()  # 地图投影信息
    im_data = dataset.ReadAsArray(0, 0, im_width, im_height).astype(np.float32)  # 将数据写成数组，对应栅格矩阵
    # im_lon=[im_geotrans[0]+i*im_geotrans[1] for i in range(im_width)]
    # im_lat=[im_geotrans[3]+i*im_geotrans[5] for i in range(im_height)]
    no_data = dataset.GetRasterBand(1).GetNoDataValue()
    
    del dataset  # 关闭对象，文件dataset   
    return im_data,im_width,im_height,no_data
#=========================
def levles_colos(p, img_out_path):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.subplots_adjust(bottom=0.85)
    fig.colorbar(p, orientation="horizontal", cax=ax)
    # ax.set_title('亿元', loc='right')
    plt.axis('off')
    # fig.show()
    plt.savefig(img_out_path, dpi=300, bbox_inches='tight', transparent=True)
    
    plt.clf()
    plt.close(fig)
#=========================
def contourf(fin,lv,cp):
    tif_file_folder,tif_file_name = os.path.split(fin)

    im_data,im_width,im_height,no_data=read_img(fin)
    im_data[im_data==no_data]=np.nan
    im_data[im_data==-9999]=np.nan
    im_data[im_data==-99]=np.nan
    # print(np.nanmin(im_data), np.nanmax(im_data))

    col = im_width
    row = im_height

    fig = plt.figure(figsize=(16, 9))
    # fig = plt.figure()
    p = plt.contourf(np.linspace(0, col, col, dtype=int),
                      np.linspace(0, row, row, dtype=int),
                      im_data[::-1],
                      levels=lv,
                      extend='both',
                      cmap=cp)

    plt.axis('off')

    # cax = fig.add_axes([0.15, 0.05, 0.7, 0.03])
    # cax.set_title('元/㎡', loc='right')
    # plt.colorbar(p, orientation="horizontal", cax=cax)

    # plt.show(bbox_inches='tight')
    plt.savefig(img_path +'\\'+tif_file_name.replace("tif", "png"), dpi=300, bbox_inches='tight', transparent=True)
    # levles_colos(p, img_path +'\\'+ tif_file_name.replace(".tif", "_color.png"))
    
    
    plt.clf()
    plt.close(fig)
    
#========================   

city=['贵阳市','六盘水','遵义市','安顺市','毕节市','铜仁市','黔西南','黔东南','黔南']

tifpath=r'H:\company\公司项目\贵州\merged\city'
tifpath2=r'H:\company\公司项目\贵州\merged'
tifpath3=r'H:\company\公司项目\贵州\merged\重采样'
img_path=r'H:\company\公司项目\贵州\img'

# for ci in city:
#     fin=tifpath+'\\总价值_'+ci+'.tif'
#     lv=np.round(np.linspace(0, 52.5, 22), 1)
#     cp="viridis"
#     contourf(fin,lv,cp)

# for ci in city:
#     fin=tifpath+'\\碳汇量2021_'+ci+'.tif'
#     lv=np.round(np.linspace(0, 1500, 16), 1)
#     cp="YlGn"
#     contourf(fin,lv,cp)

# fin=tifpath3+'\\总价值.tif'
# lv=np.round(np.linspace(0, 52.5, 22), 1)
# cp="viridis"
# contourf(fin,lv,cp)

for na in range(2010,2022):
    fin=tifpath3+'\\碳汇量'+str(na)+'.tif'
    lv=np.round(np.linspace(0, 1500, 16), 1)
    cp="YlGn"
    contourf(fin,lv,cp)
    
    for ci in city:
        fin=tifpath+'\\碳汇量'+str(na)+'_'+ci+'.tif'
        lv=np.round(np.linspace(0, 1500, 16), 1)
        cp="YlGn"
        contourf(fin,lv,cp)

    
